package com.ruoyi.framework.websocket;

import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.io.IOException;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArraySet;

@Slf4j
@Component
public class MyWebSocketHandler extends TextWebSocketHandler {

    // 存储所有活跃连接（线程安全）
    private final Set<WebSocketSession> activeSessions = new CopyOnWriteArraySet<>();

    // 存储用户ID与会话的映射（可选，用于定向消息）
    private final Map<String, WebSocketSession> userSessionMap = new ConcurrentHashMap<>();

    // 连接建立时触发
    @Override
    public void afterConnectionEstablished(WebSocketSession session) throws Exception {
        System.out.println("新连接建立：" + session.getId() + "，当前连接数：" + activeSessions.size());
        log.info("新连接建立：{}", session.getId());
        // 将新会话添加到活跃连接集合
        activeSessions.add(session);
        session.sendMessage(new TextMessage("已连接到 WebSocket（端口：" + session.getLocalAddress().getPort() + "）"));
    }

    // 接收消息时触发
    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        System.out.println("收到消息：" + message.getPayload());
        log.info("收到消息：{}", message.getPayload());
        session.sendMessage(new TextMessage("服务器已收到：" + message.getPayload()));
    }

    /**
     * 连接关闭时触发
     */
    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) throws Exception {
        // 从活跃连接中移除会话
        activeSessions.remove(session);
        System.out.println("连接关闭：" + session.getId() + "，原因：" + status.getReason() + "，当前连接数：" + activeSessions.size());

        // 从用户映射中移除
        String userId = (String) session.getAttributes().get("userId");
        if (userId != null) {
            userSessionMap.remove(userId);
        }
    }

    /**
     * 发生错误时触发
     */
    @Override
    public void handleTransportError(WebSocketSession session, Throwable exception) throws Exception {
        System.err.println("连接错误（" + session.getId() + "）：" + exception.getMessage());
        // 发送错误提示给客户端
        session.sendMessage(new TextMessage("发生错误：" + exception.getMessage()));
    }

    /**
     * 广播消息给所有连接（可选：排除发送者）
     */
    public void broadcastMessage(String message, WebSocketSession excludeSession) throws IOException {
        for (WebSocketSession session : activeSessions) {
            // 只给已打开的会话发送消息，排除指定会话
            if (session.isOpen() && !session.equals(excludeSession)) {
                session.sendMessage(new TextMessage("广播消息：" + message));
            }
        }
    }

    /**
     * 向特定用户发送消息
     */
    public void sendToUser(String userId, String message) throws IOException {
        WebSocketSession session = userSessionMap.get(userId);
        if (session != null && session.isOpen()) {
            session.sendMessage(new TextMessage("专属消息：" + message));
        } else {
            System.out.println("用户 " + userId + " 未在线或连接已关闭");
        }
    }

    /**
     * 获取当前活跃连接数
     */
    public int getActiveSessionCount() {
        return activeSessions.size();
    }
}
